#include <string>
#include <iostream>
#include <mutex>
#include "body.h"
#include "camera.h"
#include "gl_render.h"


std::string silhouette_vertex_code =
"#version 330 core\n"
"layout(location = 0) in vec3 aPos;\n"
"uniform mat4 Trans;\n"
"void main()\n"
"{\n"
"  gl_Position = Trans * vec4(aPos, 1.0);\n"
"}";

std::string silhouette_fragment_code =
"#version 330 core\n"
"uniform float SilhouetteID;\n"
"out float FragColor;\n"
"void main()\n"
"{\n"
"  FragColor = SilhouetteID;\n"
"}";

std::string normal_vertex_code =
    "#version 330 core\n"
    "layout(location = 0) in vec3 aPos;\n"
    "layout(location = 1) in vec3 aNormal;\n"
    "flat out vec3 Normal;\n"
    "uniform mat4 Trans;\n"
    "uniform mat3 Rot;\n"
    "void main()\n"
    "{\n"
    "  gl_Position = Trans * vec4(aPos, 1.0);\n"
    "  Normal = Rot * aNormal;\n"
    "}";

std::string normal_fragment_code =
    "#version 330 core\n"
    "flat in vec3 Normal;\n"
    "out vec4 FragColor;\n"
    "void main()\n"
    "{\n"
    "	 FragColor = vec4(0.5 - 0.5 * Normal, 1.0).zyxw;\n"
    "}";

static bool CheckCompileErrors(unsigned shader, const std::string& type) {
	int success;
	char info_log[1024];
	if (type != "PROGRAM") {
		glGetShaderiv(shader, GL_COMPILE_STATUS, &success);
		if (!success) {
			glGetShaderInfoLog(shader, 1024, nullptr, info_log);
			std::cerr << "Shader compilation error of type: " << type << std::endl
				<< info_log << std::endl;
			return false;
		}
	}
	else {
		glGetProgramiv(shader, GL_LINK_STATUS, &success);
		if (!success) {
			glGetProgramInfoLog(shader, 1024, nullptr, info_log);
			std::cerr << "Shader linking error of type: " << type << std::endl
				<< info_log << std::endl;
			return false;
		}
	}
	return true;
}

static bool CreateShaderProgram(
	const char* vertex_shader_code,
	const char* fragment_shader_code,
	unsigned* shader_program) 
{
	unsigned vertex_shader = glCreateShader(GL_VERTEX_SHADER);
	glShaderSource(vertex_shader, 1, &vertex_shader_code, nullptr);
	glCompileShader(vertex_shader);
	if (!CheckCompileErrors(vertex_shader, "VERTEX")) {
		return false;
	}

	unsigned fragment_shader = glCreateShader(GL_FRAGMENT_SHADER);
	glShaderSource(fragment_shader, 1, &fragment_shader_code, nullptr);
	glCompileShader(fragment_shader);
	if (!CheckCompileErrors(fragment_shader, "FRAGMENT")) {
		return false;
	}

	*shader_program = glCreateProgram();
	glAttachShader(*shader_program, vertex_shader);
	glAttachShader(*shader_program, fragment_shader);
	glLinkProgram(*shader_program);
	if (!CheckCompileErrors(*shader_program, "PROGRAM")) {
		return false;
	}

	glDeleteShader(vertex_shader);
	glDeleteShader(fragment_shader);

	return true;
}

int GLRender::n_instances_ = 0;

GLRender* GLRender::Build(const Intrinsics& intrinsics) {
	auto render_ptr = new GLRender;

	if (!glfwInit()) {
		std::cerr << "[ERROR] Failed to initialize GLFW" << std::endl;
		delete render_ptr;
		return nullptr;
	}

	glfwWindowHint(GLFW_CONTEXT_VERSION_MAJOR, 3);
	glfwWindowHint(GLFW_CONTEXT_VERSION_MINOR, 3);
	glfwWindowHint(GLFW_OPENGL_PROFILE, GLFW_OPENGL_CORE_PROFILE);
	glfwWindowHint(GLFW_RESIZABLE, GL_TRUE);
	glfwWindowHint(GLFW_VISIBLE, GL_FALSE);

	auto window = glfwCreateWindow(640, 480, "window", nullptr, nullptr);
	if (window == nullptr) {
		std::cerr << "[ERROR] Failed to create GLFW window" << std::endl;
		glfwTerminate();
		delete render_ptr;
		return nullptr;
	}

	render_ptr->intrinsics_ = intrinsics;
	render_ptr->image_width_ = intrinsics.width;
	render_ptr->image_height_ = intrinsics.height;

	render_ptr->CalculateProjectionMatrix();
	render_ptr->CalculateProjectionTerms();

	glfwMakeContextCurrent(window);

	glewExperimental = true;
	if (glewInit() != GLEW_OK) {
		std::cerr << "[ERROR] Failed to initialize GLEW" << std::endl;
		glfwDestroyWindow(window);
		window = nullptr;
		glfwTerminate();
		delete render_ptr;
		return nullptr;
	}
	render_ptr->window_ = window;

	if (!CreateShaderProgram(
		normal_vertex_code.c_str(),
		normal_fragment_code.c_str(), &render_ptr->normal_shader_))
	{
		std::cerr << "[ERROR] Create normal shader failed" << std::endl;
		delete render_ptr;
		return nullptr;
	}

	if (!CreateShaderProgram(
		silhouette_vertex_code.c_str(),
		silhouette_fragment_code.c_str(), &render_ptr->silhouette_shader_))
	{
		std::cerr << "[ERROR] Create silhouette shader failed" << std::endl;
		delete render_ptr;
		return nullptr;
	}

	render_ptr->CreateBufferObjects();

	glfwMakeContextCurrent(nullptr);

	render_ptr->ClearNormalImage();
	render_ptr->ClearDepthImage();
	render_ptr->ClearSilhouetteImage();

	n_instances_++;
	render_ptr->set_up_ = true;

	return render_ptr;
}

GLRender::~GLRender() {
	if (set_up_) DeleteBufferObjects();
	
	if (render_bodies_.size() > 0) {
		glfwMakeContextCurrent(window_);
		for (auto &render_body : render_bodies_) {
			render_body.DeleteVertexObjects();
		}
		glfwMakeContextCurrent(nullptr);
		glfwDestroyWindow(window_);
		window_ = nullptr;
		n_instances_--;
		if (n_instances_ == 0) glfwTerminate();
	}
}

void GLRender::CreateBufferObjects() {
	glGenRenderbuffers(1, &rbo_depth_);
  glBindRenderbuffer(GL_RENDERBUFFER, rbo_depth_);
  glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT16, image_width_,
                        image_height_);
  glBindRenderbuffer(GL_RENDERBUFFER, 0);

  glGenRenderbuffers(1, &rbo_normal_);
  glBindRenderbuffer(GL_RENDERBUFFER, rbo_normal_);
  glRenderbufferStorage(GL_RENDERBUFFER, GL_RGBA8, image_width_, image_height_);
  glBindRenderbuffer(GL_RENDERBUFFER, 0);

  glGenRenderbuffers(1, &rbo_silhouette_);
  glBindRenderbuffer(GL_RENDERBUFFER, rbo_silhouette_);
  glRenderbufferStorage(GL_RENDERBUFFER, GL_R8, image_width_, image_height_);
  glBindRenderbuffer(GL_RENDERBUFFER, 0);

  glGenFramebuffers(1, &fbo_);
  glBindFramebuffer(GL_FRAMEBUFFER, fbo_);
  glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,
                            GL_RENDERBUFFER, rbo_normal_);
	glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1,
                            GL_RENDERBUFFER, rbo_silhouette_);										
  glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT,
                            GL_RENDERBUFFER, rbo_depth_);
  glBindFramebuffer(GL_FRAMEBUFFER, 0);
}

void GLRender::DeleteBufferObjects() {
	glDeleteRenderbuffers(1, &rbo_silhouette_);
  glDeleteRenderbuffers(1, &rbo_normal_);
  glDeleteRenderbuffers(1, &rbo_depth_);
  glDeleteFramebuffers(1, &fbo_);
}

bool GLRender::AddBody(Body* body_ptr) {
	const std::lock_guard<std::mutex> lock{ mutex_ };

	if (nullptr == body_ptr) {
		return false;
	}

	for (auto& p : render_bodies_) {
		if (body_ptr->id() == p.body_ptr->id()) {
			std::cerr << "Body # " << body_ptr->id() 
								<< " already added to GLRender" 
								<< std::endl;
			return false;
		}
	}

	RenderBody render_body = {body_ptr};
	glfwMakeContextCurrent(window_);
	render_body.CreateVertexObjects();
	glfwMakeContextCurrent(nullptr);

	render_bodies_.push_back(std::move(render_body));

	return true;
}

bool GLRender::DeleteBody(Body* body_ptr) {
	const std::lock_guard<std::mutex> lock{ mutex_ };

	if (nullptr == body_ptr) {
		return false;
	}

	for (size_t i = 0; i < render_bodies_.size(); ++i) {
		if (body_ptr->id() == render_bodies_[i].body_ptr->id()) {
			glfwMakeContextCurrent(window_);
			render_bodies_[i].DeleteVertexObjects();
			glfwMakeContextCurrent(nullptr);

			render_bodies_.erase(std::begin(render_bodies_) + i);
			return true;
		}
	}

	std::cerr << "[ERROR] Can not delete Body # " << body_ptr->id() << std::endl;
	return false;
}

void GLRender::ClearBodies() {
  const std::lock_guard<std::mutex> lock{mutex_};
	glfwMakeContextCurrent(window_);
	for (auto &render_body : render_bodies_) {
		render_body.DeleteVertexObjects();
	}
	glfwMakeContextCurrent(nullptr);
  render_bodies_.clear();
}

bool GLRender::RenderNormal() {
  const std::lock_guard<std::mutex> lock{mutex_};

	glfwMakeContextCurrent(window_);
  bool ret = RenderNormal(projection_matrix_, world2camera_pose_);
	glfwMakeContextCurrent(nullptr);
	return ret;
}

bool GLRender::RenderSilhouette() {
  const std::lock_guard<std::mutex> lock{mutex_};

	glfwMakeContextCurrent(window_);
  bool ret = RenderSilhouette(projection_matrix_, world2camera_pose_);
	glfwMakeContextCurrent(nullptr);
	return ret;
}

void GLRender::CalculateProjectionMatrix() {
  projection_matrix_ << 2.0f * intrinsics_.fx / float(intrinsics_.width), 0.0f,
      2.0f * (intrinsics_.cx + 0.5f) / float(intrinsics_.width) - 1.0f, 0.0f,
      0.0f, 2.0f * intrinsics_.fy / float(intrinsics_.height),
      2.0f * (intrinsics_.cy + 0.5f) / float(intrinsics_.height) - 1.0f, 0.0f,
      0.0f, 0.0f, (z_max_ + z_min_) / (z_max_ - z_min_),
      -2.0f * z_max_ * z_min_ / (z_max_ - z_min_), 0.0f, 0.0f, 1.0f, 0.0f;
}

void GLRender::CalculateProjectionTerms() {
  projection_term_a_ = z_max_ * z_min_ * USHRT_MAX / (z_max_ - z_min_);
  projection_term_b_ = z_max_ * USHRT_MAX / (z_max_ - z_min_);
}

bool GLRender::RenderNormal(
    const Eigen::Matrix4f& projection_matrix,
    const Pose& world2camera_pose) {
  glViewport(0, 0, image_width_, image_height_);

  glBindFramebuffer(GL_FRAMEBUFFER, fbo_);
  glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
  glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
  glEnable(GL_DEPTH_TEST);
  glFrontFace(GL_CCW);
  glCullFace(GL_FRONT);

  glUseProgram(normal_shader_);
  for (const auto &render_data_body : render_bodies_) {
    Pose trans_without_projection{
        world2camera_pose * render_data_body.body_ptr->body2world_pose()};
    Eigen::Matrix4f trans{projection_matrix * trans_without_projection.matrix()};
    Eigen::Matrix3f rot{trans_without_projection.rotation().matrix()};

    unsigned loc;
    loc = glGetUniformLocation(normal_shader_, "Trans");
    glUniformMatrix4fv(loc, 1, GL_FALSE, trans.data());
    loc = glGetUniformLocation(normal_shader_, "Rot");
    glUniformMatrix3fv(loc, 1, GL_FALSE, rot.data());

    if (geometry_enable_culling_)
      glEnable(GL_CULL_FACE);
    else
      glDisable(GL_CULL_FACE);

    glBindVertexArray(render_data_body.vao);
    glDrawArrays(GL_TRIANGLES, 0, render_data_body.n_vertices);
    glBindVertexArray(0);
  }
  glBindFramebuffer(GL_FRAMEBUFFER, 0);

  image_rendered_ = true;
  normal_image_fetched_ = false;
  depth_image_fetched_ = false;
  return true;
}

bool GLRender::RenderSilhouette(
	const Eigen::Matrix4f& projection_matrix,
	const Pose& world2camera_pose) 
{
	glViewport(0, 0, image_width_, image_height_);

	glBindFramebuffer(GL_FRAMEBUFFER, fbo_);
	glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
	glEnable(GL_DEPTH_TEST);
	glFrontFace(GL_CCW);
	glCullFace(GL_FRONT);

	glUseProgram(silhouette_shader_);
	for (const auto& render_data_body : render_bodies_) {
		Eigen::Matrix4f trans{
				projection_matrix *
				(world2camera_pose * render_data_body.body_ptr->body2world_pose())
						.matrix() };

		unsigned loc;
		loc = glGetUniformLocation(silhouette_shader_, "Trans");
		glUniformMatrix4fv(loc, 1, GL_FALSE, trans.data());
		loc = glGetUniformLocation(silhouette_shader_, "SilhouetteID");
		glUniform1f(loc, float(render_data_body.body_ptr->id()) / 255.0f);

		if (geometry_enable_culling_)
			glEnable(GL_CULL_FACE);
		else
			glDisable(GL_CULL_FACE);

		glBindVertexArray(render_data_body.vao);
		glDrawArrays(GL_TRIANGLES, 0, render_data_body.n_vertices);
		glBindVertexArray(0);
	}
	glBindFramebuffer(GL_FRAMEBUFFER, 0);

	image_rendered_ = true;
	silhouette_image_fetched_ = false;
	depth_image_fetched_ = false;
	return true;
}

bool GLRender::FetchNormalImage(cv::Mat *normal_image) {
  if (!image_rendered_) return false;
  if (normal_image_fetched_) return true;

  glPixelStorei(GL_PACK_ALIGNMENT, (normal_image->step & 3) ? 1 : 4);
  glPixelStorei(GL_PACK_ROW_LENGTH,
                GLint(normal_image->step / normal_image->elemSize()));
  glBindFramebuffer(GL_FRAMEBUFFER, fbo_);
  glBindRenderbuffer(GL_RENDERBUFFER, rbo_normal_);
  glReadPixels(0, 0, image_width_, image_height_, GL_BGRA, GL_UNSIGNED_BYTE,
               normal_image->data);
  glBindRenderbuffer(GL_RENDERBUFFER, 0);
  glBindFramebuffer(GL_FRAMEBUFFER, 0);

  normal_image_fetched_ = true;
  return true;
}

bool GLRender::FetchSilhouetteImage(cv::Mat* silhouette_image) {
	if (!image_rendered_) return false;
	if (silhouette_image_fetched_) return true;

	glPixelStorei(GL_PACK_ALIGNMENT, (silhouette_image->step & 3) ? 1 : 4);
	glPixelStorei(GL_PACK_ROW_LENGTH,
		GLint(silhouette_image->step / silhouette_image->elemSize()));
	glBindFramebuffer(GL_FRAMEBUFFER, fbo_);
	glBindRenderbuffer(GL_RENDERBUFFER, rbo_normal_);
	glReadPixels(0, 0, image_width_, image_height_, GL_RED, GL_UNSIGNED_BYTE,
		silhouette_image->data);
	glBindRenderbuffer(GL_RENDERBUFFER, 0);
	glBindFramebuffer(GL_FRAMEBUFFER, 0);

	silhouette_image_fetched_ = true;
	return true;
}

bool GLRender::FetchDepthImage(cv::Mat* depth_image) {
	if (!image_rendered_) return false;
	if (depth_image_fetched_) return true;

	glPixelStorei(GL_PACK_ALIGNMENT, (depth_image->step & 3) ? 1 : 4);
	glPixelStorei(GL_PACK_ROW_LENGTH,
		GLint(depth_image->step / depth_image->elemSize()));
	glBindFramebuffer(GL_FRAMEBUFFER, fbo_);
	glBindRenderbuffer(GL_RENDERBUFFER, rbo_depth_);
	glReadPixels(0, 0, image_width_, image_height_, GL_DEPTH_COMPONENT,
		GL_UNSIGNED_SHORT, depth_image->data);
	glBindRenderbuffer(GL_RENDERBUFFER, 0);
	glBindFramebuffer(GL_FRAMEBUFFER, 0);

	depth_image_fetched_ = true;
	return true;
}

bool GLRender::FetchNormalImage() {
  const std::lock_guard<std::mutex> lock{mutex_};

	glfwMakeContextCurrent(window_);
  bool ret = FetchNormalImage(&normal_image_);
	glfwMakeContextCurrent(nullptr);
	return ret;
}

bool GLRender::FetchSilhouetteImage() {
  const std::lock_guard<std::mutex> lock{mutex_};

	glfwMakeContextCurrent(window_);
  bool ret = FetchSilhouetteImage(&silhouette_image_);
	glfwMakeContextCurrent(nullptr);
	return ret;
}

bool GLRender::FetchDepthImage() {
  const std::lock_guard<std::mutex> lock{mutex_};

	glfwMakeContextCurrent(window_);
  bool ret = FetchDepthImage(&depth_image_);
	glfwMakeContextCurrent(nullptr);
	return ret;
}

float GLRender::Depth(ushort depth_image_value) const {
	return projection_term_a_ / (projection_term_b_ - float(depth_image_value));
}

float GLRender::Depth(const cv::Point2i& image_coordinate) const {
	float depth_image_value = float(depth_image_.at<ushort>(image_coordinate));
	return projection_term_a_ / (projection_term_b_ - depth_image_value);
}

Eigen::Vector3f GLRender::PointVector(
	const cv::Point2i& image_coordinate) const {
	float depth_image_value = float(depth_image_.at<ushort>(image_coordinate));
	float depth = projection_term_a_ / (projection_term_b_ - depth_image_value);
	return Eigen::Vector3f{
			depth * (image_coordinate.x - intrinsics_.cx) / intrinsics_.fx,
			depth * (image_coordinate.y - intrinsics_.cy) / intrinsics_.fy,
			depth };
}

void GLRender::RenderBody::CreateVertexObjects() {
	std::vector<float> vertex_data;
	auto mesh_ptr = body_ptr->mesh();

	for (const auto& triangle_indices : mesh_ptr->indices) {
		std::array<Eigen::Vector3f, 3> points;
		for (int i = 0; i < 3; ++i) {
			auto v = mesh_ptr->vertices[triangle_indices[i]];
			points[i] = Eigen::Vector3f(v[0], v[1], v[2]);
		}

		Eigen::Vector3f normal{
				(points[2] - points[1]).cross(points[0] - points[1]).normalized() };

		for (auto point : points) {
			vertex_data.insert(end(vertex_data), point.data(), point.data() + 3);
			vertex_data.insert(end(vertex_data), normal.data(), normal.data() + 3);
		}
	}

	n_vertices = unsigned(vertex_data.size()) / 6;

  glGenVertexArrays(1, &vao);
  glBindVertexArray(vao);

  glGenBuffers(1, &vbo);
  glBindBuffer(GL_ARRAY_BUFFER, vbo);
  glBufferData(GL_ARRAY_BUFFER, vertex_data.size() * sizeof(float),
               &vertex_data.front(), GL_STATIC_DRAW);

  glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float), nullptr);
  glEnableVertexAttribArray(0);
  glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, 6 * sizeof(float),
                        (void *)(3 * sizeof(float)));
  glEnableVertexAttribArray(1);

  glBindBuffer(GL_ARRAY_BUFFER, 0);
  glBindVertexArray(0);
}

void GLRender::RenderBody::DeleteVertexObjects() {
	glDeleteBuffers(1, &vbo);
	glDeleteVertexArrays(1, &vao);
}
